import tensorflow as tf
from tensorflow import keras
import numpy as np
import os

# 1.	使用tensorflow2.0完成mnist数据集处理（每小题5分）
# (1)	数据预处理
# ①	加载mnist数据集
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
# ②	归一化处理
x_train, x_test = x_train.astype(np.float32)/255., x_test.astype(np.float32)/255.
# print(x_train.shape, x_test.shape)
# ③	相对维度进行扩张
x_train, x_test = np.expand_dims(x_train, axis=3), np.expand_dims(x_test, axis=3)
# print(x_train.shape, x_test.shape)
# ④	对标签进行独热处理
y_train, y_test = tf.one_hot(y_train, depth=10), tf.one_hot(y_test, depth=10)
# print(y_train.shape)
# (2)	跃迁块设置
# ①	自定义函数，实现3*3，步长为1卷积处理
def conv3x3(channel, kernel_size=(3, 3), strides=1):
    return keras.layers.Conv2D(channel, kernel_size, strides,
                               padding='same', kernel_initializer=tf.random_normal_initializer())
# ②	创建跃迁类，设定初始化参数
class ResNetBolock(keras.Model):

    def __init__(self, channel, strides, residual_path=False):
        super(ResNetBolock, self).__init__()

        self.channel = channel
        self.strides = strides
        self.residual_path= residual_path

        self.conv1 = conv3x3(channel, strides)
        self.bn1 = keras.layers.BatchNormalization()
        self.conv2 = conv3x3(channel)
        self.bn2 = keras.layers.BatchNormalization()

        # ③	设置跃迁处理，合理进行处理
        if residual_path:
            self.down_conv = conv3x3(channel, kernel_size=(1, 1))
            self.down_bn = keras.layers.BatchNormalization()

# ④	实现正向传播算法
    def call(self, inputs, training=None):
        x = self.bn1(inputs)
        x = self.conv1(x)
        x = keras.layers.ReLU(x)
        x = self.bn2(x)
        x = self.conv2(x)
        x = keras.layers.ReLU(x)
        # ⑤	通道数量改变情况下，进行跃迁处理
        if self.residual_path:
            # ⑥	做两次卷积以后，判断是否做1*1跃迁处理
            residual = self.down_bn(inputs)
            residual = self.down_conv(residual)
        return (x + residual)

# (3)	Resnet网络布置

# ①	创建Resnet类
class ResNet(keras.Model):
    def __init__(self, block_list, num_class, channel = 16):

        super(ResNet, self).__init__()
        # ②	定义初始化函数，设定参数值
        self.block_list = block_list
        self.num_class = num_class
        self.input_channel = channel
        self.output_channel = channel
        # 创建序列模型
        self.conv1 = conv3x3(self.input_channel)
        self.bn1 = keras.layers.BatchNormalization()
        self.block = keras.Sequential()
        self.de = ResNetBolock(self.input_channel,strides=1 )
        self.residual_path = False
# ③	进行循环判断，循环两次，通道数量翻倍
        for layers_id in self.block_list:
            for block_id in range(layers_id):
                # 不是第一层处理时，对数据进行降采样
                if layers_id != 0 and block_id==0:
                    self.de = ResNetBolock(self.output_channel, strides=2, residual_path=self.residual_path)
                else:
                    if self.input_channel != self.output_channel:
                        self.residual_path = True
                    else:
                        self.residual_path = False
                    self.de = ResNetBolock(self.output_channel, strides=1, residual_path=self.residual_path)
                self.input_channel = self.output_channel # 通道数量一致
            # 每次大循环结束，通道数量翻倍
            self.output_channel *= 2
            # 每次循环结束后，将模块进行存储
            self.block.add(self.de)
# ④	完成正向传播处理
    def call(self, inputs, training=None):
        x = inputs
        x = self.conv1(x)
        x = self.bn1(x)
        x = keras.layers.ReLU(x)
        x = self.block(x)
# ⑤	使用平均池化进行分类处理
        x = keras.layers.GlobalAveragePooling2D(x)
        x = keras.layers.Dense(10)
        return x
# (4)	主函数处理
if __name__ == '__main__':
# ①	设置合理参数
    epoch = 10
    block_list = [2, 2, 2]
    learning_rate = 0.1

# ②	创建resnet网络（2,2,2）
    model = ResNet(block_list=block_list, num_class=10)
# ③	编译网络
    model.compile(optimizer=keras.optimizers.Adam(learning_rate=learning_rate), loss=keras.losses.CategoricalCrossentropy)
    model.build([None, 28, 28, 1])
# ④	训练训练集数据
    model.fit(x_train, y_train, epochs=epoch)
# ⑤	打印最终得分
    metrics = model.evaluate(x_test, y_test, batch_size=50)
    print('最终得分为：', metrics)
#
